from typing import TypedDict, Annotated

from langchain_core.messages import HumanMessage, AIMessage
from langchain_ollama import ChatOllama
from langgraph.constants import START
from langgraph.graph import add_messages, StateGraph


class State(TypedDict):
    messages: Annotated[list, add_messages]


def chatbot(state: State):
    return {"messages": [model.invoke(state["messages"])]}


if __name__ == '__main__':
    graph_builder = StateGraph(State)
    model = ChatOllama(model="qwen3:latest", base_url="http://192.168.97.217:11434")
    graph_builder.add_node("chatbot",chatbot)
    graph_builder.add_edge(START, "chatbot")
    graph=graph_builder.compile()
    # graph_image = graph.get_graph(xray=True).draw_mermaid_png()
    # with open("graph.png", "wb") as f:
    #     f.write(graph_image)
    message_list=[
        HumanMessage(content="你好,我叫何沛，好久不见"),
        AIMessage(content="你好，我叫小智，一名乐于助人的AI助手，很高兴认识你。"),
        HumanMessage(content="还记得我叫做什么名字")
    ]
    final_state = graph.invoke({"messages": message_list})
    # final_state = graph.invoke({"messages": [{"role": "user", "content": "你好,我叫何沛，好久不见"}]})
    print(final_state["messages"][0].id)
    print(final_state["messages"][0].content)
    print("===================================")
    print(final_state["messages"][1].id)
    print(final_state["messages"][1].content)
    print("===================================")
    print(final_state)

